/*
 * Copyright 2013-2016 ksyun.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License").
 * You may not use this file except in compliance with the License.
 * A copy of the License is located at
 *
 *  http://ksyun.com/apache2.0
 *
 * or in the "license" file accompanying this file. This file is distributed
 * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
 * express or implied. See the License for the specific language governing
 * permissions and limitations under the License.
 */
package com.kingsoft.miot.bridge.iam.common.auth;


import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.UnsupportedEncodingException;
import java.net.URI;
import java.security.DigestInputStream;
import java.security.MessageDigest;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.SortedMap;
import java.util.TreeMap;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;

import org.apache.commons.codec.binary.Hex;

import com.kingsoft.miot.bridge.iam.common.KscClientException;
import com.kingsoft.miot.bridge.iam.common.SignableRequest;
import com.kingsoft.miot.bridge.iam.common.util.SdkHttpUtils;
import com.kingsoft.miot.bridge.iam.common.util.StringUtils;

import static  com.kingsoft.miot.bridge.iam.common.util.StringUtils.UTF8;

/**
 * Signer implementation that signs requests with the AWS4 signing protocol.
 */
public class AWS4Signer {

    public static final String HMAC_SHA256_ALGORITHM = "HmacSHA256";

    public static final String AWS4_TERMINATOR = "aws4_request";

    public static final String AWS4_SIGNING_ALGORITHM = "AWS4-HMAC-SHA256";

    public static final String X_AMZ_DATE = "X-Amz-Date";

    public static final String AUTHORIZATION = "Authorization";

    public static final String HOST = "Host";

    /**
     * Service name override for use when the endpoint can't be used to
     * determine the service name.
     */
    private String serviceName;

    /**
     * Region name override for use when the endpoint can't be used to determine
     * the region name.
     */
    private String regionName;

    /**
     * The scope of the signature.
     */
    private final String scope;

    /**
     * UTC formatted version of the signing time stamp.
     */
    private final String formattedSigningDateTime;

    /**
     * UTC Formatted Signing date with time stamp stripped
     */
    private final String formattedSigningDate;

    /**
     * Construct a new AWS4 signer instance. By default, enable double
     * url-encoding.
     */
    public AWS4Signer(String regionName, String serviceName) {
        long signingDateTimeMilli = System.currentTimeMillis();
        this.serviceName = serviceName;
        this.regionName  = regionName;
        this.formattedSigningDate = StringUtils.formatDateStamp(signingDateTimeMilli);
        this.formattedSigningDateTime = StringUtils.formatTimestamp(signingDateTimeMilli);
        this.scope = String.format("%s/%s/%s/%s", this.formattedSigningDate, this.regionName, this.serviceName, AWS4_TERMINATOR);
    }

    public void sign(SignableRequest<?> request, String ak, String sk) {

        addHostHeader(request);
        request.addHeader(X_AMZ_DATE, this.formattedSigningDateTime);

        String contentSha256          = calculateContentHash(request);
        final String canonicalRequest = createCanonicalRequest(request, contentSha256);
        final String stringToSign     = createStringToSign(canonicalRequest);
        final byte[] signingKey       = deriveSigningKey(sk);
        final byte[] signature        = computeSignature(stringToSign, signingKey);

        request.addHeader(AUTHORIZATION, buildAuthorizationHeader(request, signature, ak));
    }

    private void addHostHeader(SignableRequest<?> request) {
        // AWS4 requires that we sign the Host header so we
        // have to have it in the request by the time we sign.

        final URI endpoint = request.getUri();
        String hostHeader  = endpoint.getHost();
        if (SdkHttpUtils.isUsingNonDefaultPort(endpoint)) {
            hostHeader = hostHeader + ":"+endpoint.getPort();
        }
        request.addHeader(HOST, hostHeader);
    }

    /**
     * Step 1 of the AWS Signature version 4 calculation. Refer to
     * http://docs.aws
     * .amazon.com/general/latest/gr/sigv4-create-canonical-request.html to
     * generate the canonical request.
     */
    private String createCanonicalRequest(SignableRequest<?> request, String contentSha256) {
        final String strFormat = "%s\n%s\n%s\n%s\n%s\n%s";
        return String.format(strFormat, request.getHttpMethod(),
                getCanonicalizedResourcePath(request.getUri().getPath()),
                getCanonicalizedQueryString(request),
                getCanonicalizedHeaderString(request),
                getSignedHeadersString(request),
                contentSha256);
    }

    /**
     * Step 2 of the AWS Signature version 4 calculation. Refer to
     * http://docs.aws
     * .amazon.com/general/latest/gr/sigv4-create-string-to-sign.html.
     */
    private String createStringToSign(String canonicalRequest) {

        final String strFormat = "%s\n%s\n%s\n%s";

        return String.format(strFormat,
                AWS4_SIGNING_ALGORITHM,
                this.formattedSigningDateTime,
                this.scope,
                Hex.encodeHexString(hash(canonicalRequest)));
    }

    /**
     * Step 3 of the AWS Signature version 4 calculation. It involves deriving
     * the signing key and computing the signature. Refer to
     * http://docs.aws.amazon
     * .com/general/latest/gr/sigv4-calculate-signature.html
     */
    private byte[] deriveSigningKey(String sk) {

        return newSigningKey(sk,this.formattedSigningDate,
                this.regionName,
                this.serviceName);
    }

    /**
     * Step 3 of the AWS Signature version 4 calculation. It involves deriving
     * the signing key and computing the signature. Refer to
     * http://docs.aws.amazon
     * .com/general/latest/gr/sigv4-calculate-signature.html
     */
    private byte[] computeSignature(String stringToSign, byte[] signingKey) {
        try {
            return sign(stringToSign.getBytes("UTF-8"), signingKey, HMAC_SHA256_ALGORITHM);
        }catch (UnsupportedEncodingException e) {
            throw new KscClientException(e);
        }
    }

    /**
     * Creates the authorization header to be included in the request.
     */
    private String buildAuthorizationHeader(SignableRequest<?> request,
            byte[] signature, String ak) {

        final String credential      = "Credential=" +  ak + "/" + this.scope;
        final String signerHeaders   = "SignedHeaders=" + getSignedHeadersString(request);
        final String signatureHeader = "Signature=" + Hex.encodeHexString(signature);

        return String.format("%s %s, %s, %s", AWS4_SIGNING_ALGORITHM, credential, signerHeaders, signatureHeader);
    }


    private String getCanonicalizedHeaderString(SignableRequest<?> request) {
        final List<String> sortedHeaders = new ArrayList<String>(request.getHeaders()
                .keySet());
        Collections.sort(sortedHeaders, String.CASE_INSENSITIVE_ORDER);

        final Map<String, String> requestHeaders = request.getHeaders();
        StringBuilder buffer = new StringBuilder();
        for (String header : sortedHeaders) {
            String key = StringUtils.lowerCase(header);
            String value = requestHeaders.get(header);

            StringUtils.appendCompactedString(buffer, key);
            buffer.append(":");
            if (value != null) {
                StringUtils.appendCompactedString(buffer, value);
            }

            buffer.append("\n");
        }

        return buffer.toString();
    }

    private String getSignedHeadersString(SignableRequest<?> request) {
        final List<String> sortedHeaders = new ArrayList<String>(request
                .getHeaders().keySet());
        Collections.sort(sortedHeaders, String.CASE_INSENSITIVE_ORDER);

        StringBuilder buffer = new StringBuilder();
        for (String header : sortedHeaders) {
            if (buffer.length() > 0)
                buffer.append(";");
            buffer.append(StringUtils.lowerCase(header));
        }

        return buffer.toString();
    }

    /**
     * Calculate the hash of the request's payload. Subclass could override this
     * method to provide different values for "x-amz-content-sha256" header or
     * do any other necessary set-ups on the request headers. (e.g. aws-chunked
     * uses a pre-defined header value, and needs to change some headers
     * relating to content-encoding and content-length.)
     */
    private String calculateContentHash(SignableRequest<?> request){

        try {
            InputStream payloadStream = request.getContent();
            if (payloadStream == null) {
                payloadStream = new ByteArrayInputStream(new byte[0]);
            }
            String contentSha256 = Hex.encodeHexString(hash(payloadStream));
            payloadStream.reset();

            return contentSha256;
        }catch (IOException e) {
            throw new KscClientException(e);
        }
    }

    /**
     * Generates a new signing key from the given parameters and returns it.
     */

    private byte[] newSigningKey(String sk, String dateStamp, String regionName, String serviceName) {
        try {
            byte[] kSecret = ("AWS4" + sk).getBytes("UTF-8");
            byte[] kDate = sign(dateStamp, kSecret, HMAC_SHA256_ALGORITHM);
            byte[] kRegion = sign(regionName, kDate, HMAC_SHA256_ALGORITHM);
            byte[] kService = sign(serviceName, kRegion, HMAC_SHA256_ALGORITHM);

            return sign(AWS4_TERMINATOR, kService, HMAC_SHA256_ALGORITHM);
        }catch (UnsupportedEncodingException e){
            throw new KscClientException(e);
        }
    }

    private byte[] sign(String stringData, byte[] key,
                       String algorithm) throws KscClientException {
        try {
            byte[] data = stringData.getBytes(UTF8);
            return sign(data, key, algorithm);
        } catch (Exception e) {
            throw new KscClientException(
                    "Unable to calculate a request signature: "
                            + e.getMessage(), e);
        }
    }

    private byte[] sign(byte[] data, byte[] key,
                          String algorithm) throws KscClientException {
        try {
            Mac mac = Mac.getInstance(algorithm);
            mac.init(new SecretKeySpec(key, algorithm));
            return mac.doFinal(data);
        } catch (Exception e) {
            throw new KscClientException(
                    "Unable to calculate a request signature: "
                            + e.getMessage(), e);
        }
    }


    private static byte[] hash(String text) {
        try {
            MessageDigest md = MessageDigest.getInstance("SHA-256");
            md.update(text.getBytes(UTF8));
            return md.digest();
        } catch (Exception e) {
            throw new KscClientException(
                    "Unable to compute hash while signing request: "
                            + e.getMessage(), e);
        }
    }

    private byte[] hash(InputStream input) throws KscClientException {
        try {
            MessageDigest md = MessageDigest.getInstance("SHA-256");
            DigestInputStream digestInputStream = new DigestInputStream(input, md);
            byte[] buffer = new byte[1024];
            while (digestInputStream.read(buffer) > -1);
            return digestInputStream.getMessageDigest().digest();
        } catch (Exception e) {
            throw new KscClientException(
                    "Unable to compute hash while signing request: "
                            + e.getMessage(), e);
        }
    }

    private String getCanonicalizedQueryString(SignableRequest<?> request) {
        /*
         * If we're using POST and we don't have any request payload content,
         * then any request query parameters will be sent as the payload, and
         * not in the actual query string.
         */
        if (SdkHttpUtils.usePayloadForQueryParameters(request)) return "";
        return getCanonicalizedQueryString(request.getParameters());
    }

    /**
     * Examines the specified query string parameters and returns a
     * canonicalized form.
     * <p>
     * The canonicalized query string is formed by first sorting all the query
     * string parameters, then URI encoding both the key and value and then
     * joining them, in order, separating key value pairs with an '&'.
     *
     * @param parameters
     *            The query string parameters to be canonicalized.
     *
     * @return A canonicalized form for the specified query string parameters.
     */
    private String getCanonicalizedQueryString(Map<String, String> parameters) {

        final SortedMap<String, String> sorted = new TreeMap<>();

        /**
         * Signing protocol expects the param values also to be sorted after url
         * encoding in addition to sorted parameter names.
         */
        for (Map.Entry<String, String> entry : parameters.entrySet()) {
            final String encodedParamName = SdkHttpUtils.urlEncode(entry.getKey(), false);
            final String encodedValues    = SdkHttpUtils.urlEncode(entry.getValue(), false);
            sorted.put(encodedParamName, encodedValues);
        }

        final StringBuilder result = new StringBuilder();
        for(Map.Entry<String, String> entry : sorted.entrySet()) {
            if (result.length() > 0) {
                result.append("&");
            }
            result.append(entry.getKey())
                    .append("=")
                    .append(entry.getValue());
        }

        return result.toString();
    }

    private String getCanonicalizedResourcePath(String resourcePath) {
        if (resourcePath == null || resourcePath.isEmpty()) {
            return "/";
        }

        return resourcePath.startsWith("/") ? resourcePath :  "/".concat(resourcePath);
    }
}
